package com.seo.textgen.markov;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Random;

import com.seo.textgen.tokenizer.Token;

public class MapMarkovModel implements MarkovModel {

	private Random random = new Random(5);
	
	private Map<Chain, Collection<Token>> model = new HashMap<Chain, Collection<Token>>();

	public MapMarkovModel() {
	}
	
	public void addChain(Chain subChain, Token token) {
		Collection<Token> probableTokens = model.get(subChain);
		if (probableTokens == null) {
			probableTokens = new ArrayList<Token>();
			model.put(subChain, probableTokens);
		}
		probableTokens.add(token);
	}

	public Token getNextToken(Chain chain) {
		Collection<Token> tokens = model.get(chain);
		if (tokens == null) {
			return null;
		}
		float rand = random.nextFloat();
		float sum = 0;
		Iterator<Token> it = tokens.iterator();
		
		Token token = null;
		while (it.hasNext()) {
			token = it.next();
			sum += token.getProbability();
			if (sum > rand) {
				return token;
			}
		}
		return token;
	}
	
}
